Write a custom CUDA kernel to optimize `AOAF` (Adaptive Offset Activation Function).

Formula: f(x) = max(0, x - 0.17 * alpha) + 0.17 * alpha
Where `alpha` is the mean of the input feature tensor.

Problem Analysis:
1. Dynamic Parameter: The activation function's behavior depends on the mean of the input tensor, which must be computed on-the-fly for each forward pass.
2. Memory Bottleneck: A standard PyTorch implementation requires a full reduction pass to compute the mean, followed by several element-wise operations (mul, sub, max, add), leading to multiple global memory accesses.

Optimization Strategy: Fused Two-Pass Kernel

1. One-Block-per-Row (or Sample): Launch one CUDA block for each row of the input tensor, assuming the mean is computed row-wise.

2. Fused Two-Pass Algorithm:
   - Pass 1 (Statistics): Threads within a block cooperatively compute the sum of their assigned row using a parallel reduction in Shared Memory. Thread 0 calculates the mean.
   - Pass 2 (Apply): After the mean is broadcasted via Shared Memory, threads iterate over their portion of the row again. They read the input value, apply the AOAF formula `max(0, x - c) + c` (where `c` is the pre-computed offset), and write the result.

3. Vectorization: Use `float4` for all memory accesses to maximize bandwidth.
  
Here's an example to show you the syntax of inline embedding custom CUDA operators in torch: The example given architecture is:   
  
```python
import torch
import torch.nn as nn

# --- 基准测试配置 ---
# 大尺寸 Tensor
BATCH_SIZE = 4096
HIDDEN_DIM = 4096
SHAPE = (BATCH_SIZE, HIDDEN_DIM)

class AOAF(nn.Module):
    """
    An Adaptive Offset Activation Function for CNN Image Classification Tasks
    https://www.mdpi.com/2079-9292/11/22/3799.
    Formula: f(x) = max(0, x - 0.17 * alpha) + 0.17 * alpha
    Where `alpha` is the mean of the input feature tensor.
    """
    def __init__(self):
        super(AOAF, self).__init__()
        self.beta1 = 0.17
        self.beta2 = 0.17

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        alpha = x.mean(dim=-1, keepdim=True)
        
        offset = self.beta1 * alpha
        
        return torch.clamp(x - offset, min=0.0) + self.beta2 * alpha

class Model(nn.Module):
    def __init__(self):
        super(Model, self).__init__()
        self.act = AOAF()
    
    def forward(self, x):
        return self.act(x)

def get_inputs():
    input_tensor = torch.randn(SHAPE, dtype=torch.float32) * 5.0
    return [input_tensor.contiguous()]

def get_init_inputs():
    return []